{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp train_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# train_utils\n",
    "\n",
    "> Basic utils for experiments - create databanch, learner with fastai v1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "# from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai.basic_train import *\n",
    "from fastai.vision import *\n",
    "from fastai.script import *\n",
    "from fastai.vision.models.xresnet import *\n",
    "# from model_constructor.net import *\n",
    "# from model_constructor.layers import SimpleSelfAttention, ConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import math\n",
    "import torch\n",
    "from torch.optim.optimizer import Optimizer, required\n",
    "import itertools as it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mish & Ranger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Mish(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "#         print(\"Mish activation loaded...\")\n",
    "\n",
    "    def forward(self, x):  \n",
    "        #save 1 second per epoch with no x= x*() and then return x...just inline it.\n",
    "        return x *( torch.tanh(F.softplus(x))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "#Ranger deep learning optimizer - RAdam + Lookahead combined.\n",
    "  #https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer\n",
    "\n",
    "  #Ranger has now been used to capture 12 records on the FastAI leaderboard.\n",
    "\n",
    "  #This version = 9.3.19  \n",
    "\n",
    "  #Credits:\n",
    "  #RAdam -->  https://github.com/LiyuanLucasLiu/RAdam\n",
    "  #Lookahead --> rewritten by lessw2020, but big thanks to Github @LonePatient and @RWightman for ideas from their code.\n",
    "  #Lookahead paper --> MZhang,G Hinton  https://arxiv.org/abs/1907.08610\n",
    "\n",
    "  #summary of changes: \n",
    "  #full code integration with all updates at param level instead of group, moves slow weights into state dict (from generic weights), \n",
    "  #supports group learning rates (thanks @SHolderbach), fixes sporadic load from saved model issues.\n",
    "  #changes 8/31/19 - fix references to *self*.N_sma_threshold; \n",
    "                  #changed eps to 1e-5 as better default than 1e-8.\n",
    "\n",
    "class Ranger(Optimizer):\n",
    "\n",
    "    def __init__(self, params, lr=1e-3, alpha=0.5, k=6, N_sma_threshhold=5, betas=(.95,0.999), eps=1e-5, weight_decay=0):\n",
    "        #parameter checks\n",
    "        if not 0.0 <= alpha <= 1.0:\n",
    "            raise ValueError(f'Invalid slow update rate: {alpha}')\n",
    "        if not 1 <= k:\n",
    "            raise ValueError(f'Invalid lookahead steps: {k}')\n",
    "        if not lr > 0:\n",
    "            raise ValueError(f'Invalid Learning Rate: {lr}')\n",
    "        if not eps > 0:\n",
    "            raise ValueError(f'Invalid eps: {eps}')\n",
    "\n",
    "        #parameter comments:\n",
    "        # beta1 (momentum) of .95 seems to work better than .90...\n",
    "        #N_sma_threshold of 5 seems better in testing than 4.\n",
    "        #In both cases, worth testing on your dataset (.90 vs .95, 4 vs 5) to make sure which works best for you.\n",
    "\n",
    "        #prep defaults and init torch.optim base\n",
    "        defaults = dict(lr=lr, alpha=alpha, k=k, step_counter=0, betas=betas, N_sma_threshhold=N_sma_threshhold, eps=eps, weight_decay=weight_decay)\n",
    "        super().__init__(params,defaults)\n",
    "\n",
    "        #adjustable threshold\n",
    "        self.N_sma_threshhold = N_sma_threshhold\n",
    "\n",
    "        #now we can get to work...\n",
    "        #removed as we now use step from RAdam...no need for duplicate step counting\n",
    "        #for group in self.param_groups:\n",
    "        #    group[\"step_counter\"] = 0\n",
    "            #print(\"group step counter init\")\n",
    "\n",
    "        #look ahead params\n",
    "        self.alpha = alpha\n",
    "        self.k = k \n",
    "\n",
    "        #radam buffer for state\n",
    "        self.radam_buffer = [[None,None,None] for ind in range(10)]\n",
    "\n",
    "        #self.first_run_check=0\n",
    "\n",
    "        #lookahead weights\n",
    "        #9/2/19 - lookahead param tensors have been moved to state storage.  \n",
    "        #This should resolve issues with load/save where weights were left in GPU memory from first load, slowing down future runs.\n",
    "\n",
    "        #self.slow_weights = [[p.clone().detach() for p in group['params']]\n",
    "        #                     for group in self.param_groups]\n",
    "\n",
    "        #don't use grad for lookahead weights\n",
    "        #for w in it.chain(*self.slow_weights):\n",
    "        #    w.requires_grad = False\n",
    "\n",
    "    def __setstate__(self, state):\n",
    "        print(\"set state called\")\n",
    "        super(Ranger, self).__setstate__(state)\n",
    "\n",
    "\n",
    "    def step(self, closure=None):\n",
    "        loss = None\n",
    "        #note - below is commented out b/c I have other work that passes back the loss as a float, and thus not a callable closure.  \n",
    "        #Uncomment if you need to use the actual closure...\n",
    "\n",
    "        #if closure is not None:\n",
    "            #loss = closure()\n",
    "\n",
    "        #Evaluate averages and grad, update param tensors\n",
    "        for group in self.param_groups:\n",
    "\n",
    "            for p in group['params']:\n",
    "                if p.grad is None:\n",
    "                    continue\n",
    "                grad = p.grad.data.float()\n",
    "                if grad.is_sparse:\n",
    "                    raise RuntimeError('Ranger optimizer does not support sparse gradients')\n",
    "\n",
    "                p_data_fp32 = p.data.float()\n",
    "\n",
    "                state = self.state[p]  #get state dict for this param\n",
    "\n",
    "                if len(state) == 0:   #if first time to run...init dictionary with our desired entries\n",
    "                    #if self.first_run_check==0:\n",
    "                        #self.first_run_check=1\n",
    "                        #print(\"Initializing slow buffer...should not see this at load from saved model!\")\n",
    "                    state['step'] = 0\n",
    "                    state['exp_avg'] = torch.zeros_like(p_data_fp32)\n",
    "                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n",
    "\n",
    "                    #look ahead weight storage now in state dict \n",
    "                    state['slow_buffer'] = torch.empty_like(p.data)\n",
    "                    state['slow_buffer'].copy_(p.data)\n",
    "\n",
    "                else:\n",
    "                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n",
    "                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)\n",
    "\n",
    "                #begin computations \n",
    "                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n",
    "                beta1, beta2 = group['betas']\n",
    "\n",
    "                #compute variance mov avg\n",
    "                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n",
    "                #compute mean moving avg\n",
    "                exp_avg.mul_(beta1).add_(1 - beta1, grad)\n",
    "\n",
    "                state['step'] += 1\n",
    "\n",
    "\n",
    "                buffered = self.radam_buffer[int(state['step'] % 10)]\n",
    "                if state['step'] == buffered[0]:\n",
    "                    N_sma, step_size = buffered[1], buffered[2]\n",
    "                else:\n",
    "                    buffered[0] = state['step']\n",
    "                    beta2_t = beta2 ** state['step']\n",
    "                    N_sma_max = 2 / (1 - beta2) - 1\n",
    "                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)\n",
    "                    buffered[1] = N_sma\n",
    "                    if N_sma > self.N_sma_threshhold:\n",
    "                        step_size = math.sqrt((1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / (N_sma_max - 2)) / (1 - beta1 ** state['step'])\n",
    "                    else:\n",
    "                        step_size = 1.0 / (1 - beta1 ** state['step'])\n",
    "                    buffered[2] = step_size\n",
    "\n",
    "                if group['weight_decay'] != 0:\n",
    "                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)\n",
    "\n",
    "                if N_sma > self.N_sma_threshhold:\n",
    "                    denom = exp_avg_sq.sqrt().add_(group['eps'])\n",
    "                    p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom)\n",
    "                else:\n",
    "                    p_data_fp32.add_(-step_size * group['lr'], exp_avg)\n",
    "\n",
    "                p.data.copy_(p_data_fp32)\n",
    "\n",
    "                #integrated look ahead...\n",
    "                #we do it at the param level instead of group level\n",
    "                if state['step'] % group['k'] == 0:\n",
    "                    slow_p = state['slow_buffer'] #get access to slow param tensor\n",
    "                    slow_p.add_(self.alpha, p.data - slow_p)  #(fast weights - slow weights) * alpha\n",
    "                    p.data.copy_(slow_p)  #copy interpolated weights to RAdam param tensor\n",
    "\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_data(size=128, woof=1, bs=64, workers=None, **kwargs):\n",
    "    if woof:\n",
    "        path = URLs.IMAGEWOOF    # if woof \n",
    "    else:\n",
    "        path = URLs.IMAGENETTE\n",
    "    path = untar_data(path)\n",
    "    print('data path  ', path)\n",
    "    n_gpus = num_distrib() or 1\n",
    "    if workers is None: workers = min(8, num_cpus()//n_gpus)\n",
    "    return (ImageList.from_folder(path).split_by_folder(valid='val')\n",
    "            .label_from_folder().transform(([flip_lr(p=0.5)], []), size=size)\n",
    "            .databunch(bs=bs, num_workers=workers)\n",
    "            .presize(size, scale=(0.35,1))\n",
    "            .normalize(imagenet_stats))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data path   /root/.fastai/data/imagewoof2\n"
     ]
    }
   ],
   "source": [
    "data = get_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ImageDataBunch;\n",
       "\n",
       "Train: LabelList (9025 items)\n",
       "x: ImageList\n",
       "Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128)\n",
       "y: CategoryList\n",
       "n02099601,n02099601,n02099601,n02099601,n02099601\n",
       "Path: /root/.fastai/data/imagewoof2;\n",
       "\n",
       "Valid: LabelList (3929 items)\n",
       "x: ImageList\n",
       "Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128),Image (3, 128, 128)\n",
       "y: CategoryList\n",
       "n02099601,n02099601,n02099601,n02099601,n02099601\n",
       "Path: /root/.fastai/data/imagewoof2;\n",
       "\n",
       "Test: None"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#hide\n",
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_learn(\n",
    "        gpu:Param(\"GPU to run on\", str)=None,\n",
    "        woof: Param(\"Use imagewoof (otherwise imagenette)\", int)=1,\n",
    "        size: Param(\"Size (px: 128,192,224)\", int)=128,\n",
    "        alpha: Param(\"Alpha\", float)=0.99, \n",
    "        mom: Param(\"Momentum\", float)=0.95, #? 0.9\n",
    "        eps: Param(\"epsilon\", float)=1e-6,\n",
    "        bs: Param(\"Batch size\", int)=64,\n",
    "        mixup: Param(\"Mixup\", float)=0.,\n",
    "        opt: Param(\"Optimizer (adam,rms,sgd)\", str)='ranger',\n",
    "        sa: Param(\"Self-attention\", int)=0,\n",
    "        sym: Param(\"Symmetry for self-attention\", int)=0,\n",
    "        model: Param('model as partial', callable) = xresnet50\n",
    "        ):\n",
    " \n",
    "    if   opt=='adam' : opt_func = partial(optim.Adam, betas=(mom,alpha), eps=eps)\n",
    "    elif opt=='ranger'  : opt_func = partial(Ranger,  betas=(mom,alpha), eps=eps)\n",
    "    data = get_data(size, woof, bs)\n",
    "    learn = (Learner(data, model(), wd=1e-2, opt_func=opt_func,\n",
    "             metrics=[accuracy,top_k_accuracy],\n",
    "             bn_wd=False, true_wd=True,\n",
    "             loss_func = LabelSmoothingCrossEntropy(),))\n",
    "    print('Learn path', learn.path)\n",
    "    if mixup: learn = learn.mixup(alpha=mixup)\n",
    "    return learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = get_learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.get_learn(gpu: Param(help='GPU to run on', type=<class 'str'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = None, woof: Param(help='Use imagewoof (otherwise imagenette)', type=<class 'int'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 1, size: Param(help='Size (px: 128,192,224)', type=<class 'int'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 128, alpha: Param(help='Alpha', type=<class 'float'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 0.99, mom: Param(help='Momentum', type=<class 'float'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 0.95, eps: Param(help='epsilon', type=<class 'float'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 1e-06, bs: Param(help='Batch size', type=<class 'int'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 64, mixup: Param(help='Mixup', type=<class 'float'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 0.0, opt: Param(help='Optimizer (adam,rms,sgd)', type=<class 'str'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 'ranger', sa: Param(help='Self-attention', type=<class 'int'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 0, sym: Param(help='Symmetry for self-attention', type=<class 'int'>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = 0, model: Param(help='model as partial', type=<built-in function callable>, opt=True, action=None, nargs=None, const=None, choices=None, required=None) = functools.partial(<function xresnet at 0x7f870702a3b0>, expansion=4, n_layers=[3, 4, 6, 3], name='xresnet50'))>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#hide\n",
    "learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
